import xml.etree.ElementTree as ET
from collections import defaultdict
import aiohttp
import asyncio
from tqdm.asyncio import tqdm as tqdm_asyncio
from datetime import datetime, timezone, timedelta
import gzip
import shutil
from xml.dom import minidom
import re
from opencc import OpenCC
import os
from tqdm import tqdm
from io import BytesIO
from typing import Dict, List, Tuple, Optional

# =========================
# 全局与工具
# =========================

RE_SPACES = re.compile(r"\s+")
TZ_CST8 = timezone(timedelta(hours=8))  # +0800

_CC: Optional[OpenCC] = None
def get_opencc() -> Optional[OpenCC]:
    global _CC
    if _CC is None:
        try:
            _CC = OpenCC("t2s")  # 繁体转简体
        except Exception as e:
            print(f"OpenCC 初始化失败，将跳过繁简转换: {e}")
            _CC = None
    return _CC

def zh_hans(text: Optional[str]) -> str:
    if not text:
        return ""
    cc = get_opencc()
    if not cc:
        return text
    try:
        return cc.convert(text)
    except Exception as e:
        print(f"convert to zh_hans failed: {e}")
        return text

def normalize_name(name: str) -> str:
    # 去空格后作为 channel id 使用（与原逻辑兼容）
    return RE_SPACES.sub("", name or "")

def ensure_dir(path: str):
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)

def is_gzip_bytes(data: bytes) -> bool:
    # gzip 文件魔数：1f 8b
    return len(data) >= 2 and data[0] == 0x1F and data[1] == 0x8B

def parse_epg_datetime_to_cst8(dt_str: str) -> str:
    """
    将 EPG 的 start/stop 时间解析并转换到 +0800 时区，返回 "YYYYmmddHHMMSS %z" 格式。
    输入可为:
      - "YYYYmmddHHMMSS +0800"
      - "YYYYmmddHHMMSS+0800"
      - "YYYYmmddHHMMSS" (无时区，视为 +0800)
    """
    if not dt_str:
        # 兜底返回当前时间
        return datetime.now(TZ_CST8).strftime("%Y%m%d%H%M%S %z")

    s = RE_SPACES.sub("", dt_str)
    # 尝试带时区
    dt = None
    for fmt in ("%Y%m%d%H%M%S%z", "%Y%m%d%H%M%S"):
        try:
            dt = datetime.strptime(s, fmt)
            break
        except Exception:
            continue

    if dt is None:
        # 解析失败时兜底
        return datetime.now(TZ_CST8).strftime("%Y%m%d%H%M%S %z")

    # 如果无时区信息，视为 +0800
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=TZ_CST8)
    else:
        # 转到 +0800
        dt = dt.astimezone(TZ_CST8)
    return dt.strftime("%Y%m%d%H%M%S %z")

# =========================
# 网络获取与解压
# =========================

async def fetch_one(session: aiohttp.ClientSession, url: str) -> Optional[bytes]:
    try:
        async with session.get(url) as resp:
            if resp.status != 200:
                print(f"{url} 返回状态码: {resp.status}")
                return None
            data = await resp.read()

            # 判断是否 gzip（注意：与 Content-Encoding 不同，此处是文件本身为 gzip）
            ctype = resp.headers.get("Content-Type", "").lower()
            is_gz = url.lower().endswith(".gz") or ("gzip" in ctype) or is_gzip_bytes(data)

            if is_gz:
                try:
                    data = gzip.decompress(data)
                except Exception as e:
                    print(f"{url} 解压 gzip 失败: {e}")
                    return None

            # 返回 XML 字节（ElementTree 可直接解析 bytes 并识别编码）
            return data
    except asyncio.TimeoutError:
        print(f"{url} 请求超时")
    except aiohttp.ClientError as e:
        print(f"{url} HTTP 请求错误: {e}")
    except Exception as e:
        print(f"{url} 其他错误: {e}")
    return None

# =========================
# EPG 解析与合并
# =========================

def parse_epg(epg_bytes: bytes) -> Tuple[Dict[str, str], Dict[str, List[ET.Element]]]:
    """
    输入为 XML 字节（若来源为 .gz，已在 fetch 阶段解压）
    返回:
      - channels: {channel_id(normalized): display_name(human)}
      - programmes: {channel_id(normalized): [programme_elements]}
    """
    try:
        root = ET.fromstring(epg_bytes)  # 自动识别 XML 声明里的编码
    except ET.ParseError as e:
        print(f"XML 解析失败: {e}")
        snippet = epg_bytes[:500]
        try:
            print(f"问题片段(utf-8 尝试解码): {snippet.decode('utf-8','ignore')}")
        except Exception:
            print("问题片段无法解码为 utf-8")
        return {}, defaultdict(list)

    # 先建立 id -> (normalized_id, display_name)
    id_to_norm_and_display: Dict[str, Tuple[str, str]] = {}
    for channel in root.findall("channel"):
        cid_raw = channel.get("id") or ""
        cid = zh_hans(cid_raw)
        # 找 display-name（可能有多个，取第一个）
        dn_elem = channel.find("display-name")
        dn_txt = zh_hans(dn_elem.text) if dn_elem is not None else ""
        if not dn_txt:
            dn_txt = cid
        norm_id = normalize_name(dn_txt)
        id_to_norm_and_display[cid] = (norm_id, dn_txt)

    channels: Dict[str, str] = {}  # norm_id -> display_name
    programmes: Dict[str, List[ET.Element]] = defaultdict(list)

    for prog in root.findall("programme"):
        cid_raw = prog.get("channel") or ""
        cid = zh_hans(cid_raw)
        norm_id: str
        display_name: str

        if cid in id_to_norm_and_display:
            norm_id, display_name = id_to_norm_and_display[cid]
        else:
            # 未在 channel 清单中出现过，回退使用原 cid
            display_name = zh_hans(cid) or "Unknown"
            norm_id = normalize_name(display_name)

        # 处理时间并转换到 +0800
        start_raw = prog.get("start") or ""
        stop_raw = prog.get("stop") or ""
        start_fmt = parse_epg_datetime_to_cst8(start_raw)
        stop_fmt = parse_epg_datetime_to_cst8(stop_raw)

        # 构造新的 programme 元素（避免在原树上新增节点）
        new_prog = ET.Element("programme", attrib={
            "channel": norm_id,
            "start": start_fmt,
            "stop": stop_fmt
        })

        # 标题与简介
        title_elem = prog.find("title")
        title_txt = zh_hans(title_elem.text) if title_elem is not None else ""
        new_title = ET.SubElement(new_prog, "title")
        new_title.text = title_txt

        desc_elem = prog.find("desc")
        if desc_elem is not None and (desc_elem.text is not None):
            desc_txt = zh_hans(desc_elem.text)
            new_desc = ET.SubElement(new_prog, "desc")
            new_desc.text = desc_txt

        # 保存
        if norm_id not in channels:
            channels[norm_id] = display_name
        programmes[norm_id].append(new_prog)

    return channels, programmes

def write_to_xml(channels: Dict[str, str], programmes: Dict[str, List[ET.Element]],
                 filename: str, pretty: bool = False) -> None:
    ensure_dir(os.path.dirname(filename) or ".")
    current_time = datetime.now(TZ_CST8).strftime("%Y%m%d%H%M%S %z")
    root = ET.Element("tv", attrib={"date": current_time})

    # 写入 channel 清单
    for channel_id, display_name in channels.items():
        ch = ET.SubElement(root, "channel", attrib={"id": channel_id})
        dn = ET.SubElement(ch, "display-name", attrib={"lang": "zh"})
        dn.text = display_name

    # 写入 programme
    for channel_id, plist in programmes.items():
        for p in plist:
            # 确保 channel 属性是我们最终的 channel_id
            p.set("channel", channel_id)
            root.append(p)

    if pretty:
        # 漂亮打印（较慢）
        rough_string = ET.tostring(root, encoding="utf-8")
        reparsed = minidom.parseString(rough_string)
        with open(filename, "w", encoding="utf-8") as f:
            f.write(reparsed.toprettyxml(indent="\t", newl="\n"))
    else:
        # 高性能输出
        tree = ET.ElementTree(root)
        tree.write(filename, encoding="utf-8", xml_declaration=True)

def compress_to_gz(input_filename: str, output_filename: str, compresslevel: int = 6) -> None:
    ensure_dir(os.path.dirname(output_filename) or ".")
    with open(input_filename, "rb") as f_in:
        with gzip.open(output_filename, "wb", compresslevel=compresslevel) as f_out:
            shutil.copyfileobj(f_in, f_out)

# =========================
# 配置读取
# =========================

def get_urls(cfg_file: str = "config.txt") -> list:
    urls = []
    if not os.path.exists(cfg_file):
        print(f"未找到配置文件: {cfg_file}")
        return urls
    with open(cfg_file, "r", encoding="utf-8") as file:
        for line in file:
            line = line.strip()
            if line and not line.startswith("#"):
                urls.append(line)
    return urls

# =========================
# 主流程
# =========================

async def main():
    urls = get_urls()
    if not urls:
        print("config.txt 为空或不存在，请配置要抓取的 EPG URL")
        return

    timeout = aiohttp.ClientTimeout(total=60, connect=15)  # 可按需调整
    connector = aiohttp.TCPConnector(limit=32, ssl=False)  # 并发连接数可按网络情况调整
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/113.0.0.0 Safari/537.36"
    }

    print("Fetching EPG data...")
    async with aiohttp.ClientSession(connector=connector, timeout=timeout, headers=headers, trust_env=True) as session:
        tasks = [fetch_one(session, url) for url in urls]
        epg_bytes_list = await tqdm_asyncio.gather(*tasks, desc="Fetching URLs")

    # 解析与合并
    all_channels: Dict[str, str] = {}                   # norm_id -> display_name
    all_programmes: Dict[str, List[ET.Element]] = defaultdict(list)

    print("Parsing EPG data...")
    with tqdm(total=len(epg_bytes_list), desc="Parsing EPG", unit="file") as pbar:
        for content in epg_bytes_list:
            if not content:
                pbar.update(1)
                continue

            channels, programmes = parse_epg(content)

            # 合并：以 norm_id 为 key，display_name 以最先出现的为准
            for cid, dname in channels.items():
                if cid not in all_channels:
                    all_channels[cid] = dname

            for cid, plist in programmes.items():
                all_programmes[cid].extend(plist)

            pbar.update(1)

    # 写文件
    print("Writing to XML...")
    ensure_dir("output")
    xml_path = "output/epg.xml"
    write_to_xml(all_channels, all_programmes, xml_path, pretty=False)  # 默认高性能输出

    # 同时压缩为 .gz
    gz_path = "output/epg.gz"
    compress_to_gz(xml_path, gz_path)
    print(f"Done. 写入: {xml_path}，压缩为: {gz_path}")

if __name__ == "__main__":
    asyncio.run(main())